def assemble(assembly_code):
    register_map = {
        'ACC': 0,
        'BR': 1,
        'MR': 2,
        'AX': 3,
        'BX': 4,
        'CX': 5,
        'DX': 6,
        'SP': 7
    }

    opcodes = {
        'LOAD': [{'type': 'direct_address', 'opcode':1}],
        'STORE': [{'type': 'direct_address', 'opcode':2}],
        'MOV': [
            {'type': 'reg_reg', 'opcode':3},
            {'type': 'reg_imm', 'opcode':4}  # 新增reg_imm类型
        ],
        'ADD': [
            {'type': 'reg', 'opcode':5},
            {'type': 'immediate', 'opcode':6}
        ],
        'SUB': [
            {'type': 'reg', 'opcode':7},
            {'type': 'immediate', 'opcode':8}
        ],
        'MUL': [
            {'type': 'reg', 'opcode':9},
            {'type': 'immediate', 'opcode':10}
        ],
        'AND': [
            {'type': 'reg', 'opcode':11},
            {'type': 'immediate', 'opcode':12}
        ],
        'OR': [
            {'type': 'reg', 'opcode':13},
            {'type': 'immediate', 'opcode':14}
        ],
        'NOT': [{'type': 'none', 'opcode':15}],
        'SHL': [
            {'type': 'reg', 'opcode':16},
            {'type': 'immediate', 'opcode':17}
        ],
        'SHR': [
            {'type': 'reg', 'opcode':18},
            {'type': 'immediate', 'opcode':19}
        ],
        'CMP': [
            {'type': 'reg', 'opcode':20},
            {'type': 'immediate', 'opcode':21}
        ],
        'JMP': [{'type': 'direct_address', 'opcode':22}],
        'JZ': [{'type': 'direct_address', 'opcode':23}],
        'JNZ': [{'type': 'direct_address', 'opcode':24}],
        'JG': [{'type': 'direct_address', 'opcode':25}],
        'JNG': [{'type': 'direct_address', 'opcode':26}],
        'CALL': [{'type': 'direct_address', 'opcode':27}],
        'RET': [{'type': 'none', 'opcode':28}],
        'HALT': [{'type': 'none', 'opcode':29}],
        'PUSH': [{'type': 'reg', 'opcode':30}],
        'POP': [{'type': 'reg', 'opcode':31}]
    }

    machine_code = []

    for line in assembly_code.split('\n'):
        line = line.strip()
        if not line or line.startswith(';'):
            continue
        index = line.find(";")
        line=line[:index] if index != -1 else line
        parts = line.split()
        instruction = parts[0].upper()
        operands = ' '.join(parts[1:]) if len(parts) > 1 else ''

        # 确定操作数类型
        determined_type = None
        if instruction == 'MOV' and ',' in operands:
            split_ops = [op.strip().upper() for op in operands.split(',', 1)]
            if len(split_ops) == 2:
                if split_ops[1].startswith('#'):
                    determined_type = 'reg_imm'
                elif all(op in register_map for op in split_ops):
                    determined_type = 'reg_reg'
                else:
                    raise ValueError(f"Invalid MOV operands: {operands}")
        if not determined_type:
            if operands.startswith('#'):
                determined_type = 'immediate'
            elif ',' in operands:
                split_ops = [op.strip().upper() for op in operands.split(',')]
                if len(split_ops) == 2 and all(op in register_map for op in split_ops):
                    determined_type = 'reg_reg'
                else:
                    raise ValueError(f"Invalid operands for {instruction}: {operands}")
            else:
                if operands.strip().upper() in register_map:
                    determined_type = 'reg'
                else:
                    if operands == '':
                        determined_type = 'none'
                    else:
                        determined_type = 'direct_address'

        # 查找对应指令条目
        opcode_info = None
        for entry in opcodes.get(instruction, []):
            if entry['type'] == determined_type:
                opcode_info = entry
                break
        if not opcode_info:
            raise ValueError(f"Instruction {instruction} does not support operand type {determined_type}")

        opcode = opcode_info['opcode']
        instr_type = opcode_info['type']

        # 生成操作码和操作数二进制
        opcode_bin = format(opcode, '06b')
        operand_bin = ''

        if instr_type == 'direct_address':
            addr = int(operands)
            if not 0 <= addr <= 255:
                raise ValueError(f"Direct address {addr} out of range (0-255)")
            operand_bin = format(addr, '08b').zfill(10)
        elif instr_type == 'reg_reg':
            rd, rs = [op.strip().upper() for op in operands.split(',')]
            rd_num = register_map[rd]
            rs_num = register_map[rs]
            operand_bin = f"{rd_num:03b}{rs_num:03b}0000"
        elif instr_type == 'reg':
            rd = operands.strip().upper()
            rd_num = register_map[rd]
            operand_bin = f"{rd_num:03b}0000000"
        elif instr_type == 'immediate':
            imm = int(operands[1:])
            if not 0 <= imm <= 1023:
                raise ValueError(f"Immediate value {imm} out of range (0-1023)")
            operand_bin = format(imm, '010b')
        elif instr_type == 'reg_imm':  # 新增reg_imm处理
            rd_str, imm_str = [op.strip().upper() for op in operands.split(',', 1)]
            rd_num = register_map.get(rd_str, -1)
            imm = int(imm_str[1:])
            if rd_num == -1:
                raise ValueError(f"Invalid register: {rd_str}")
            if not 0 <= imm <= 127:
                raise ValueError(f"Immediate value {imm} out of range (0-127)")
            operand_bin = f"{rd_num:03b}{imm:07b}"
        elif instr_type == 'none':
            operand_bin = '0' * 10

        full_instruction = opcode_bin + operand_bin
        machine_code.append(full_instruction)

    return machine_code

# 示例用法
assembly = """
MOV SP,#20
MOV BX,#0
MOV AX,#1
MOV ACC,BX
ADD AX
MOV BX,ACC
MOV ACC,AX
ADD #1
MOV AX,ACC
CMP #100
JNG 3
MOV ACC,BX
HALT        ; 直接寻址
"""
machine_code=assemble(assembly)
print(machine_code)  # 输出对应的机器码列表        

with open('instruction.coe', 'w', newline='', encoding='utf-8') as outfile_coe:
    outfile_coe.write("MEMORY_INITIALIZATION_RADIX=2;\n")
    outfile_coe.write("MEMORY_INITIALIZATION_VECTOR=\n")
    for i, row in enumerate(machine_code):
        outfile_coe.write(row)
        if i < len(machine_code) - 1:
            outfile_coe.write(",\n")
        else:
            outfile_coe.write(";\n")